//
//  Microfacet.hpp
//  RayTracer
//
//  Created by 俞云康 on 4/6/16.
//  Copyright © 2016 yushroom. All rights reserved.
//

#ifndef Microfacet_hpp
#define Microfacet_hpp

#include "BxDF.hpp"

class MicrofacetDistribution {
private:
    
public:
    virtual ~MicrofacetDistribution() = default;
    virtual float D(const Direction& wh) const = 0;
    virtual void sample_f(const Direction& wo, Direction* wi, const Sample& sample, float*pdf) const = 0;
    virtual float pdf(const Direction& wo, const Direction& wi) const = 0;
};

class Blinn : public MicrofacetDistribution {
private:
    float exponent;
public:
    Blinn(float e) {
        if (e > 10000.f || isnan(e)) {
            warning("Blinn::Blinn(e), invald e\n");
            e = 10000.f;
        }
        exponent = e;
    }
    virtual float D(const Direction& wh) const override {
        float cos_theta_h = AbsCosTheta(wh);
        return (exponent+2) * INV_TWOPI * powf(cos_theta_h, exponent);
    }
    
    virtual void sample_f(const Direction& wo, Direction* wi, const Sample& sample, float*pdf) const override {
        // Compute sampled half-angle vector $\wh$ for Blinn distribution
        float costheta = powf(sample.x, 1.f / (exponent+1));
        float sintheta = sqrtf(std::max(0.f, 1.f - costheta*costheta));
        float phi = sample.y * 2.f * M_PI;
        Direction wh = SphericalDirection(sintheta, costheta, phi);
        if (!SameHemisphere(wo, wh)) wh = -wh;
        
        // Compute incident direction by reflecting about $\wh$
        *wi = -wo + 2.f * Dot(wo, wh) * wh;
        
        // Compute PDF for $\wi$ from Blinn distribution
		float blinn_pdf = 0.f;
        if (Dot(wo, wh) > 0.f) 
			blinn_pdf = ((exponent + 1.f) * powf(costheta, exponent)) / (2.f * M_PI * 4.f * Dot(wo, wh));
        *pdf = blinn_pdf;
    }
    
    virtual float pdf(const Direction& wo, const Direction& wi) const override {
        Direction wh = Normalize(wo + wi);
        float costheta = AbsCosTheta(wh);
        // Compute PDF for $\wi$ from Blinn distribution
		float blinn_pdf = 0.f;
        if (Dot(wo, wh) > 0.f) 
			blinn_pdf = ((exponent + 1.f) * powf(costheta, exponent)) / (2.f * M_PI * 4.f * Dot(wo, wh));
        return blinn_pdf;
    }
};

class Microfacet : public BxDF {
private:
    Color R;
    shared_ptr<MicrofacetDistribution> distribution;
    shared_ptr<Fresnel> fresnel;
    //Fresnel fresnel;
public:
    Microfacet(const HitInfo* hit, const Color& reflectance, shared_ptr<Fresnel> f, shared_ptr<MicrofacetDistribution> d)
    : BxDF(hit, BxDFType(BxDF_REFLECTION | BxDF_GLOSSY)), R(reflectance), distribution(d), fresnel(f) {
    }
    
    virtual Color f(const Direction& woL, const Direction& wiL) const override
    {
        float cos_theta_o = AbsCosTheta(woL);
        float cos_theta_i = AbsCosTheta(wiL);
        if (cos_theta_i == 0.f || cos_theta_o == 0.f) {
            return Color(0, 0, 0);
        }
        Direction wh = wiL + woL;
        if (wh.x == 0. && wh.y == 0.f && wh.z == 0.) {
            return Color(0, 0, 0);
        }
        wh.normalize_self();
        float cos_theta_h = Dot(wiL, wh);
        Color F = fresnel->Evaluate(cos_theta_h);
        Assert(!isnan(distribution->D(wh)));
        return R * distribution->D(wh) * G(woL, wiL, wh) * F / (4.0f * cos_theta_i * cos_theta_o);
    }
    
    virtual Color sample_f(const Direction& woL, Direction* wiL, const Sample& sample, float*pdf) const override {
        distribution->sample_f(woL, wiL, sample, pdf);
        if (!SameHemisphere(woL, *wiL)) {
            return Color(0, 0, 0);
        }
        return f(woL, *wiL);
    }
    
    virtual float pdf(const Direction& woL, const Direction& wiL) const override {
        if (!SameHemisphere(woL, wiL))
            return 0.f;
        return distribution->pdf(woL, wiL);
    }
    
    float G(const Direction& wo, const Direction& wi, const Direction& wh) const {
        float n_dot_wh = AbsCosTheta(wh);
        float n_dot_wo = AbsCosTheta(wo);
        float n_dot_wi = AbsCosTheta(wi);
        float wo_dot_wh = AbsDot(wo, wh);
        return std::min(1.f, std::min(2.f * n_dot_wh * n_dot_wo / wo_dot_wh,
                                      2.f * n_dot_wh * n_dot_wi / wo_dot_wh));
    }
};

#endif /* Microfacet_hpp */
